#include "hybrid_server.h"

int check_prikey_file()
{
// 检查私钥文件是否存在
#ifdef _WIN32
  if ((0 == _access("./KEY/kyberPriKey.bin", 0)) && (0 == access("./KEY/SMPriKey.bin", F_OK)))
    return 1;
  else
    return 0;
#else
  if ((0 == access("./KEY/kyberPriKey.bin", F_OK)) && (0 == access("./KEY/SMPriKey.bin", F_OK)))
    return 1;
  else
    return 0;
#endif
}

int gene_key()
{
  int tmp;

  // Kyber
  uint8_t pk_gen[KYBER_INDCPA_PUBLICKEYBYTES]; // 生成的公钥
  uint8_t sk_gen[KYBER_INDCPA_PUBLICKEYBYTES]; // 生成的私钥

  FILE *outKyberPriFile; // 保存私钥的文件指针
  FILE *outKyberPubFile; // 保存公钥的文件指针

  // 函数时间
  struct timeval li_gene_start, li_gene_end;
  double gene_time;

  // Generete Kyber Key

  // 获取开始时间
  gettimeofday(&li_gene_start,NULL);

  indcpa_keypair(pk_gen, sk_gen);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_gene_end,NULL);
	gene_time = 1000000*(li_gene_end.tv_sec-li_gene_start.tv_sec)+(li_gene_end.tv_usec-li_gene_start.tv_usec);
	printf("\nKyber密钥生成时间: %.0fus\n", gene_time);

  // save secret key
  if (!(outKyberPriFile = fopen("./KEY/kyberPriKey.bin", "wb")))
  {
    return -1;
  }
  fwrite(sk_gen, sizeof(char), KYBER_INDCPA_SECRETKEYBYTES, outKyberPriFile);
  fclose(outKyberPriFile);
  printf("Save Kyber SK Success\n");

  // save public key
  if (!(outKyberPubFile = fopen("./KEY/kyberPubKey.bin", "wb")))
  {
    return -1;
  }
  fwrite(pk_gen, sizeof(char), KYBER_INDCPA_PUBLICKEYBYTES, outKyberPubFile);
  fclose(outKyberPubFile);
  printf("Save Kyber PK Success\n");
  // End Kyber

  // SM2
  SM2_KEY sm2_key;    // SM2公私钥结构
  FILE *outSMPriFile; // 保存私钥的文件指针
  FILE *outSMPubFile; // 保存公钥的文件指针

  // 获取开始时间
  gettimeofday(&li_gene_start,NULL);

  sm2_key_generate(&sm2_key);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_gene_end,NULL);
	gene_time = 1000000*(li_gene_end.tv_sec-li_gene_start.tv_sec)+(li_gene_end.tv_usec-li_gene_start.tv_usec);
	printf("\nSM2密钥生成时间: %.0fus\n", gene_time);

  // save secret key
  if (!(outSMPriFile = fopen("./KEY/SMPriKey.bin", "wb")))
  {
    return -1;
  }
  fwrite(sm2_key.private_key, sizeof(char), 32, outSMPriFile);
  fclose(outSMPriFile);
  printf("Save SM2 SK Success\n");

  // save public key
  if (!(outSMPubFile = fopen("./KEY/SMPubKey.bin", "wb")))
  {
    return -1;
  }
  fwrite(&sm2_key.public_key, sizeof(SM2_POINT), 1, outSMPubFile);
  fclose(outSMPubFile);
  printf("Save SM2 PK Success\n");
  // End SM2

  return 1;
}

int hybrid_decrypt(unsigned char *cipherText, int cipherLen, unsigned char *message)
{
  if (0 == check_prikey_file())
    return -1;

  unsigned char ctKyber[KYBER_INDCPA_BYTES] = {0}; // Kyber要处理的密文
  unsigned char ctSM[SM2_CIPHER_LEN] = {0};        // SM2要处理的密文
  FILE *inKyberPriFile;                            // Kyber读取私钥的文件指针
  FILE *inSMPriFile;                               // SM2读取私钥的文件指针

  // 函数时间
  struct timeval li_dec_start, li_dec_end;
  double dec_time;

  // Kyber
  unsigned char kyberPriKey[KYBER_INDCPA_PUBLICKEYBYTES]; // Kyber的私钥

  memcpy(ctKyber, cipherText, KYBER_INDCPA_BYTES);

  // read public key
  if (!(inKyberPriFile = fopen("./KEY/kyberPriKey.bin", "rb")))
    return -1;
  fread(kyberPriKey, sizeof(char), KYBER_INDCPA_SECRETKEYBYTES, inKyberPriFile);
  fclose(inKyberPriFile);

  // 获取开始时间
  gettimeofday(&li_dec_start,NULL);

  indcpa_dec(message, ctKyber, kyberPriKey);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_dec_end,NULL);
	dec_time = 1000000*(li_dec_end.tv_sec-li_dec_start.tv_sec)+(li_dec_end.tv_usec-li_dec_start.tv_usec);
	printf("\nKyber解密时间: %.0fus\n", dec_time);
  // End Kyber

  // SM2
  SM2_KEY SMPriKey;       // SM2的私钥
  size_t SMInLen;         // SM2密文长度
  unsigned char ptSM[32]; // SM解密后的明文信息
  size_t SMOutLen;        // SM2解密后明文长度

  if (!(inSMPriFile = fopen("./KEY/SMPriKey.bin", "rb")))
    return -1;
  fread(SMPriKey.private_key, sizeof(char), 32, inSMPriFile);
  fclose(inSMPriFile);

  memcpy(ctSM, cipherText + KYBER_INDCPA_BYTES, SM2_CIPHER_LEN);
  SMInLen = cipherLen - KYBER_INDCPA_BYTES;

  // 获取开始时间
  gettimeofday(&li_dec_start,NULL);

  sm2_decrypt(&SMPriKey, ctSM, SMInLen, ptSM, &SMOutLen);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_dec_end,NULL);
	dec_time = 1000000*(li_dec_end.tv_sec-li_dec_start.tv_sec)+(li_dec_end.tv_usec-li_dec_start.tv_usec);
	printf("\nSM2解密时间: %.0fus\n", dec_time);

  memcpy(message + ENCRYPT_KEY_LEN, ptSM, ENCRYPT_KEY_LEN);
  // End SM2

  return 1;
}

int cal_sm3(unsigned char *key, unsigned char *sm3hash)
{
  SM3_CTX sm3_ctx; // SM3结构

  sm3_init(&sm3_ctx);
  sm3_update(&sm3_ctx, key, ORIGIN_KEY_LEN);
  sm3_finish(&sm3_ctx, sm3hash);
}

int main()
{
  unsigned char originHash[HASH_LEN];                                 // 原始密钥的哈希，客户端传入，长度32
  unsigned char cipherKey[KYBER_INDCPA_BYTES + SM2_CIPHER_LEN] = {0}; // 密文，客户端传入
  int cipherLen = 0;                                                  // 密文长度，客户端传入
  unsigned char unpackKey[ORIGIN_KEY_LEN] = {0};                      // 解密后密钥，本地计算
  unsigned char unpackHash[HASH_LEN];

  gene_key();
  printf("密钥文件已生成，请将公钥文件发送给用户\n");

  int listenfd, connfd;
  struct sockaddr_in servaddr;
  char socket_buff[MAXLINE];
  int n;

  if ((listenfd = socket(AF_INET, SOCK_STREAM, 0)) == -1)
  {
    printf("创建Socket套接字失败: %s(errno: %d)\n", strerror(errno), errno);
    exit(0);
  }

  memset(&servaddr, 0, sizeof(servaddr));
  servaddr.sin_family = AF_INET;
  servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
  servaddr.sin_port = htons(7658);

  if (bind(listenfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) == -1)
  {
    printf("绑定Socket套接字失败: %s(errno: %d)\n", strerror(errno), errno);
    exit(0);
  }

  if (listen(listenfd, 10) == -1)
  {
    printf("监听Socket套接字失败: %s(errno: %d)\n", strerror(errno), errno);
    exit(0);
  }

  printf("======等待用户连接======\n");
  while (1)
  {
    if ((connfd = accept(listenfd, (struct sockaddr *)NULL, NULL)) == -1)
    {
      printf("接收连接请求失败: %s(errno: %d)", strerror(errno), errno);
      continue;
    }
    n = recv(connfd, socket_buff, MAXLINE, 0);
    socket_buff[n] = '\0';

    // cipherLen||originHash||cipherKey
    // 密文长度||原始密钥哈希值||密文s
    memcpy(&cipherLen, socket_buff, sizeof(int));
    memcpy(originHash, socket_buff + sizeof(int), HASH_LEN);
    memcpy(cipherKey, socket_buff + sizeof(int) + HASH_LEN, cipherLen);

    hybrid_decrypt(cipherKey, cipherLen, unpackKey);
    printf("解密密钥(Kyber||SM2):\n");
    for (int i = 0; i < ORIGIN_KEY_LEN; i++)
    {
      printf("%02X", unpackKey[i]);
      if (i == 32)
        printf(" ");
    }
    printf("\n");

    cal_sm3(unpackKey, unpackHash);
    printf("解密密钥Hash:\n");
    for (int i = 0; i < HASH_LEN; i++)
    {
      printf("%02X", unpackHash[i]);
    }
    printf("\n");

    if (memcmp(originHash, unpackHash, HASH_LEN) != 0)
      printf("密钥验证失败\n\n");
    else
      printf("密钥验证成功\n\n");

    close(connfd);
  }

  close(listenfd);

  return 1;
}